package org.ktx.osr.audio

import android.media.MediaCodec
import android.media.MediaFormat
import android.media.MediaMuxer
import android.util.Log
import java.io.File

/**
 * 音频编码器
 * 负责将PCM音频数据编码为AAC格式
 */
class AudioEncoder(
    private val outputFile: File,
    private val sampleRate: Int = 44100,
    private val channelCount: Int = 2
) {
    private var encoder: MediaCodec? = null
    private var muxer: MediaMuxer? = null
    private var audioTrackIndex: Int = -1
    private var isMuxerStarted = false
    private var isInitialized = false

    companion object {
        private const val TAG = "AudioEncoder"
        private const val BIT_RATE = 128000
        private const val AAC_PROFILE = 2 // AAC LC profile
        private const val MAX_INPUT_SIZE = 16384
    }

    /**
     * 初始化编码器
     */
    fun initialize(): Boolean {
        if (isInitialized) return true
        
        return try {
            // 创建MediaMuxer
            muxer = MediaMuxer(outputFile.absolutePath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
            
            // 创建AAC编码器
            encoder = MediaCodec.createEncoderByType("audio/mp4a-latm")
            
            val format = MediaFormat.createAudioFormat("audio/mp4a-latm", sampleRate, channelCount).apply {
                setInteger(MediaFormat.KEY_BIT_RATE, BIT_RATE)
                setInteger(MediaFormat.KEY_AAC_PROFILE, AAC_PROFILE)
                setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, MAX_INPUT_SIZE)
                setInteger(MediaFormat.KEY_SAMPLE_RATE, sampleRate)
                setInteger(MediaFormat.KEY_CHANNEL_COUNT, channelCount)
            }
            
            Log.d(TAG, "Configuring AAC encoder with format: $format")
            
            encoder?.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            encoder?.start()
            
            isInitialized = true
            Log.d(TAG, "AAC encoder initialized successfully")
            true
        } catch (e: Exception) {
            Log.e(TAG, "Failed to initialize AAC encoder", e)
            release()
            false
        }
    }

    /**
     * 编码PCM数据
     */
    fun encodePcmData(pcmData: ByteArray, dataSize: Int) {
        if (!isInitialized) return
        
        try {
            val encoder = this.encoder ?: return
            
            // 输入数据到编码器
            val inputBufferIndex = encoder.dequeueInputBuffer(0)
            if (inputBufferIndex >= 0) {
                val inputBuffer = encoder.getInputBuffer(inputBufferIndex)
                inputBuffer?.clear()
                inputBuffer?.put(pcmData, 0, dataSize)
                encoder.queueInputBuffer(inputBufferIndex, 0, dataSize, 0, 0)
            }
            
            // 处理编码输出
            processOutput()
        } catch (e: Exception) {
            Log.e(TAG, "Error encoding PCM data", e)
        }
    }

    /**
     * 处理编码器输出
     */
    private fun processOutput() {
        try {
            val encoder = this.encoder ?: return
            val muxer = this.muxer ?: return
            val bufferInfo = MediaCodec.BufferInfo()
            var outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0)
            
            while (outputBufferIndex >= 0) {
                val outputBuffer = encoder.getOutputBuffer(outputBufferIndex)
                
                // 检查是否需要添加轨道
                if (audioTrackIndex == -1 && (bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
                    val format = encoder.outputFormat
                    audioTrackIndex = muxer.addTrack(format)
                    muxer.start()
                    isMuxerStarted = true
                    Log.d(TAG, "Added audio track to muxer: $audioTrackIndex")
                }
                
                if (outputBuffer != null && bufferInfo.size > 0 && isMuxerStarted) {
                    muxer.writeSampleData(audioTrackIndex, outputBuffer, bufferInfo)
                }
                
                encoder.releaseOutputBuffer(outputBufferIndex, false)
                outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0)
            }
        } catch (e: Exception) {
            Log.e(TAG, "Error processing encoder output", e)
        }
    }

    /**
     * 完成编码
     */
    fun finishEncoding() {
        if (!isInitialized) return
        
        try {
            val encoder = this.encoder ?: return
            val muxer = this.muxer ?: return
            
            // 发送END_OF_STREAM标志
            val inputBufferIndex = encoder.dequeueInputBuffer(10000)
            if (inputBufferIndex >= 0) {
                encoder.queueInputBuffer(
                    inputBufferIndex, 0, 0, 0,
                    MediaCodec.BUFFER_FLAG_END_OF_STREAM
                )
            }
            
            // 处理剩余的编码数据
            val bufferInfo = MediaCodec.BufferInfo()
            var outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 10000)
            
            while (outputBufferIndex >= 0) {
                val outputBuffer = encoder.getOutputBuffer(outputBufferIndex)
                
                // 检查是否需要添加轨道（如果还没有添加）
                if (audioTrackIndex == -1 && (bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
                    val format = encoder.outputFormat
                    audioTrackIndex = muxer.addTrack(format)
                    muxer.start()
                    isMuxerStarted = true
                    Log.d(TAG, "Added audio track to muxer: $audioTrackIndex")
                }
                
                if (outputBuffer != null && bufferInfo.size > 0 && isMuxerStarted) {
                    muxer.writeSampleData(audioTrackIndex, outputBuffer, bufferInfo)
                }
                
                val isEndOfStream = (bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0
                encoder.releaseOutputBuffer(outputBufferIndex, false)
                
                if (isEndOfStream) break
                
                outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 10000)
            }
            
            // 停止并释放MediaMuxer
            if (isMuxerStarted) {
                muxer.stop()
            }
            muxer.release()
            
            Log.d(TAG, "AAC encoding finished successfully")
        } catch (e: Exception) {
            Log.e(TAG, "Error finishing AAC encoding", e)
        }
    }

    /**
     * 释放资源
     */
    fun release() {
        try {
            encoder?.stop()
            encoder?.release()
            muxer?.release()
        } catch (e: Exception) {
            Log.e(TAG, "Error releasing encoder resources", e)
        } finally {
            encoder = null
            muxer = null
            audioTrackIndex = -1
            isMuxerStarted = false
            isInitialized = false
        }
    }
}
